# !/usr/bin/python3
# -*- coding:utf-8 -*-
"""
@author: JHC000abc@gmail.com
@file: db_utils.py
@time: 2024/10/24 18:16 
@desc: 

"""
import traceback
from sqlalchemy.pool import NullPool
from sqlalchemy import create_engine, exc, text
from sqlalchemy.orm import sessionmaker
from db.base_model import Base
from sqlalchemy.exc import IntegrityError
from config import DB_CONFIG, SQLPAGESIZE


class SqlalchemyDB(object):
    """
    SqlalchemyDB
    """
    _instance = None

    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self, username=None, password=None, host=None, port=None, db=None):
        if username and password and host and port and db:
            self.username = username
            self.password = password
            self.host = host
            self.port = port
            self.db = db
        else:
            self.username = DB_CONFIG["username"]
            self.password = DB_CONFIG["password"]
            self.host = DB_CONFIG["host"]
            self.port = DB_CONFIG["port"]
            self.db = DB_CONFIG["db"]

        self.db_url = f'mysql+pymysql://{self.username}:{self.password}@{self.host}:{self.port}/{self.db}'
        # print("self.db_url", self.db_url)
        self.engine = None  # 初始化引擎为 None
        self.session_maker = None  # 初始化会话工厂为 None

    def create_engine(self):
        """创建数据库引擎并建表"""
        if self.engine is None:  # 只在没有引擎时创建
            try:
                # 创建数据库引擎，使用连接池参数配置
                self.engine = create_engine(
                    self.db_url,
                    pool_recycle=3600,
                    echo=False,
                    pool_pre_ping=True,
                    poolclass=NullPool
                )
                # 建表
                Base.metadata.create_all(self.engine)
            except exc.SQLAlchemyError as e:
                print(f"Error creating SQLAlchemy engine: {e}")
                raise

    def get_session(self):
        """创建Session类，用于数据库会话管理"""
        if self.session_maker is None:  # 只在没有会话工厂时创建
            self.create_engine()  # 确保引擎已创建
            self.session_maker = sessionmaker(bind=self.engine)
        return self.session_maker()  # 返回一个会话实例

    def select(self, model=None, conditions=None, sql=None, offset=None):
        """

        :param model:
        :param conditions:{}
        :param sql:
        :param limit:
        :param offset:
        :return:
        """
        result = None
        with self.get_session() as session:
            try:
                if sql:
                    stmt = text(sql)
                    result = session.execute(stmt).fetchall()
                else:
                    query = session.query(model)
                    all_len = query.count()
                    # 根据条件过滤
                    query = query.filter_by(**conditions).limit(SQLPAGESIZE)
                    if offset or offset == 0:
                        query = query.offset(SQLPAGESIZE * offset)
                        result = {
                            "result": query.all(),
                            "all_nums": all_len,
                            "page": offset,
                            "page_size": SQLPAGESIZE
                        }
                    else:
                        result = query.all()
            except:
                print(traceback.print_exc())
            finally:
                if session:
                    session.close()
        return result

    def update(self, model=None, conditions=None, sql=None, batch=False):
        """

        :param model:
        :param conditions:
        :param sql:
        :param batch:
        :return:
        """
        result = 0
        if conditions:
            if len(conditions) > 10:
                batch = True
        with self.get_session() as session:
            try:
                if sql:
                    stmt = text(sql)
                    res = session.execute(stmt)
                    result = res.rowcount
                else:
                    tasks = []
                    for condition in conditions:
                        tasks.append(model(**condition))
                    if not batch:
                        for task in tasks:
                            session.merge(task)
                    else:
                        session.add_all(tasks)
                    result = len(tasks)
                session.commit()
            except IntegrityError:
                print("已存在", conditions)
                result = 0
                session.rollback()
            except Exception:
                print(traceback.print_exc())
                session.rollback()
                result = 0
            finally:
                if session:
                    session.close()
        return result

    def insert(self, model=None, conditions=None, batch=False):
        """
        list
        :param model:
        :param conditions:
        :param batch:
        :return:
        """
        result = None
        # 根据数据长度自动判断是否采用批量插入
        if conditions:
            if len(conditions) > 10:
                batch = True
        with self.get_session() as session:
            try:
                if batch:
                    session.bulk_insert_mappings(model, conditions)
                else:
                    tasks = []
                    for condition in conditions:
                        tasks.append(model(**condition))
                    session.add_all(tasks)
                    result = tasks
                session.commit()
            except IntegrityError:
                print("已存在", conditions)
                session.rollback()
                result = None
            except Exception:
                print("插入失败", conditions)
                session.rollback()
                result = None
            finally:
                if session:
                    session.close()
        return result
